"""
FastAPI应用入口

提供HTTP接口，支持：
- 金相组织分析（分割分析）
- 球化组织分析（分类分析）
- 带状组织分析（分类+分割分析）
"""

import os
import sys
import json
from typing import List, Optional
from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel

# 添加项目根目录到Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(current_dir)
sys.path.insert(0, project_root)

from core.jinxiang_ZZ import SegmentationAnalyzerService
from core.jinxiang_BALL import BallLevelService
from core.jinxiang_FF import DeterminationOfBandedStructureOfSteel_Service
from config import JinxiangZZ, JinxiangBall, JinxiangFF, API, MODEL_PATH, DEFAULT_INPUT_FOLDER, DEFAULT_OUTPUT_FOLDER, DEFAULT_TARGET_CLASSES, DEFAULT_DEVICE

# 创建FastAPI应用
app = FastAPI(
    title=API.title,
    description=API.description,
    version=API.version
)

# 服务实例缓存，避免重复初始化模型
_service_cache = {
    'zz': {},      # key: (model_path, device)
    'ball': {},    # key: (model_path, device)
    'ff': {}       # key: (class_model_path, seg_model_path, class_device, seg_device)
}


class AnalyzeRequest(BaseModel):
    """分析请求模型"""
    input_folder: str
    output_folder: str
    model_path: Optional[str] = None
    target_classes: Optional[List[int]] = None
    device: Optional[str] = None


class FFAnalyzeRequest(BaseModel):
    """带状组织分析请求模型"""
    input_folder: str
    output_folder: str
    class_model_path: Optional[str] = None
    seg_model_path: Optional[str] = None
    class_device: Optional[str] = None
    seg_device: Optional[str] = None


class AnalyzeResponse(BaseModel):
    """分析响应模型"""
    success: bool
    results:str
    message: str
    total_images: int = 0
    processed_images: int = 0
    failed_images: int = 0
    output_folder: str = ""


@app.get("/")
async def root():
    """根路径"""
    return {
        "message": "金相组织分析API服务正在运行",
        "endpoints": {
            "jinxiang_zz": "/api/v1/jinxiang_zz/analyze",
            "jinxiang_ball": "/api/v1/jinxiang_ball/analyze",
            "jinxiang_ff": "/api/v1/jinxiang_ff/analyze"
        }
    }


@app.get("/health")
async def health_check():
    """健康检查"""
    return {"status": "healthy", "message": "服务正常运行"}


@app.get("/config")
async def get_config():
    """获取配置信息"""
    return {
        "jinxiang_zz": {
            "model_path": JinxiangZZ.model_path,
            "input_folder": JinxiangZZ.input_folder,
            "output_folder": JinxiangZZ.output_folder,
            "target_classes": JinxiangZZ.target_classes,
            "device": JinxiangZZ.device
        },
        "jinxiang_ball": {
            "model_path": JinxiangBall.model_path,
            "input_folder": JinxiangBall.input_folder,
            "output_folder": JinxiangBall.output_folder,
            "device": JinxiangBall.device
        },
        "jinxiang_ff": {
            "class_model_path": JinxiangFF.class_model_path,
            "seg_model_path": JinxiangFF.seg_model_path,
            "input_folder": JinxiangFF.input_folder,
            "output_folder": JinxiangFF.output_folder,
            "class_device": JinxiangFF.class_device,
            "seg_device": JinxiangFF.seg_device
        },
        "api": {
            "host": API.host,
            "port": API.port,
            "title": API.title,
            "description": API.description,
            "version": API.version
        }
    }


@app.post("/api/v1/jinxiang_zz/analyze", response_model=AnalyzeResponse)
async def jinxiang_zz_analyze(request: AnalyzeRequest):
    """
    金相组织分析接口
    
    分析文件夹中的所有图像，进行金相组织的分割分析。
    
    Args:
        request: 分析请求参数
            - input_folder: 输入文件夹路径（必填）
            - output_folder: 输出文件夹路径（必填）
            - model_path: 模型路径（可选，默认使用config.py中的值）
            - target_classes: 目标类别列表（可选，默认使用config.py中的值）
            - device: 推理设备（可选，默认使用config.py中的值）
        
    Returns:
        AnalyzeResponse: 分析结果
    """
    try:
        # 使用配置中的默认值或用户提供的值
        model_path = request.model_path or MODEL_PATH
        input_folder = request.input_folder
        output_folder = request.output_folder
        target_classes = request.target_classes or DEFAULT_TARGET_CLASSES
        device = request.device or DEFAULT_DEVICE
        
        # 使用缓存的服务实例，避免重复初始化模型
        cache_key = (model_path, device)
        if cache_key not in _service_cache['zz']:
            _service_cache['zz'][cache_key] = SegmentationAnalyzerService(
                model_path=model_path,
                device=device
            )
        service = _service_cache['zz'][cache_key]
        
        # 执行分析
        result = service.analyze_folder(
            input_folder=input_folder,
            output_folder=output_folder,
            target_classes=target_classes
        )
        print ("http return:",AnalyzeResponse(
            success=result['success'],
            results=json.dumps(result['results']),
            message=result['message'],
            total_images=result['total_images'],
            processed_images=result['processed_images'],
            failed_images=result['failed_images'],
            output_folder=output_folder
        ))
        return AnalyzeResponse(
            success=result['success'],
            results=json.dumps(result['results']),
            message=result['message'],
            total_images=result['total_images'],
            processed_images=result['processed_images'],
            failed_images=result['failed_images'],
            output_folder=output_folder
        )
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"分析过程中发生错误: {str(e)}")


@app.post("/api/v1/jinxiang_ball/analyze", response_model=AnalyzeResponse)
async def jinxiang_ball_analyze(request: AnalyzeRequest):
    """
    球化组织分析接口
    
    分析文件夹中的所有图像，进行球化组织的分类分析。
    
    Args:
        request: 分析请求参数
            - input_folder: 输入文件夹路径（必填）
            - output_folder: 输出文件夹路径（必填）
            - model_path: 模型路径（可选，默认使用config.py中的值）
            - device: 推理设备（可选，默认使用config.py中的值）
        
    Returns:
        AnalyzeResponse: 分析结果
    """
    try:
        # 使用配置中的默认值或用户提供的值
        model_path = request.model_path or JinxiangBall.model_path
        input_folder = request.input_folder
        output_folder = request.output_folder
        device = request.device or JinxiangBall.device
        
        # 使用缓存的服务实例，避免重复初始化模型
        cache_key = (model_path, device)
        if cache_key not in _service_cache['ball']:
            _service_cache['ball'][cache_key] = BallLevelService(
                model_path=model_path,
                device=device
            )
        service = _service_cache['ball'][cache_key]
        
        # 执行分析
        result = service.analyze_folder(
            input_folder=input_folder,
            output_folder=output_folder
        )
        print ("http return:",AnalyzeResponse(
            success=result['success'],
            results=json.dumps(result['results']),
            message=result['message'],
            total_images=result['total_images'],
            processed_images=result['processed_images'],
            failed_images=result['failed_images'],
            output_folder=output_folder
        ))
        return AnalyzeResponse(
            success=result['success'],
            results=json.dumps(result['results']),
            message=result['message'],
            total_images=result['total_images'],
            processed_images=result['processed_images'],
            failed_images=result['failed_images'],
            output_folder=output_folder
        )
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"分析过程中发生错误: {str(e)}")


@app.post("/api/v1/jinxiang_ff/analyze", response_model=AnalyzeResponse)
async def jinxiang_ff_analyze(request: FFAnalyzeRequest):
    """
    带状组织分析接口
    
    分析文件夹中的所有图像，进行带状组织的分类和分割分析。
    
    Args:
        request: 分析请求参数
            - input_folder: 输入文件夹路径（必填）
            - output_folder: 输出文件夹路径（必填）
            - class_model_path: 分类模型路径（可选，默认使用config.py中的值）
            - seg_model_path: 分割模型路径（可选，默认使用config.py中的值）
            - class_device: 分类推理设备（可选，默认使用config.py中的值）
            - seg_device: 分割推理设备（可选，默认使用config.py中的值）
        
    Returns:
        AnalyzeResponse: 分析结果
    """
    try:
        # 使用配置中的默认值或用户提供的值
        class_model_path = request.class_model_path or JinxiangFF.class_model_path
        seg_model_path = request.seg_model_path or JinxiangFF.seg_model_path
        input_folder = request.input_folder
        output_folder = request.output_folder
        class_device = request.class_device or JinxiangFF.class_device
        seg_device = request.seg_device or JinxiangFF.seg_device
        
        # 使用缓存的服务实例，避免重复初始化模型
        cache_key = (class_model_path, seg_model_path, class_device, seg_device)
        if cache_key not in _service_cache['ff']:
            _service_cache['ff'][cache_key] = DeterminationOfBandedStructureOfSteel_Service(
                classmodel_path=class_model_path,
                segmodel_path=seg_model_path,
                classdevice=class_device,
                segdevice=seg_device
            )
        service = _service_cache['ff'][cache_key]
        
        # 执行分析
        result = service.analyze_folder(
            imgsDir_path=input_folder,
            outResultsDir_path=output_folder
        )
        
        # 转换返回格式
        # total_images = len(result) if result else 0
        # processed_images = len([r for r in result if r and ',' in r]) if result_list else 0
        # failed_images = total_images - processed_images
        print ("http return:",AnalyzeResponse(
            success=result['success'],
            results=json.dumps(result['results']),
            message=result['message'],
            total_images=result['total_images'],
            processed_images=result['processed_images'],
            failed_images=result['failed_images'],
            output_folder=output_folder
        ))
        return AnalyzeResponse(
            success=result['success'],
            results=json.dumps(result['results']),
            message=result['message'],
            total_images=result['total_images'],
            processed_images=result['processed_images'],
            failed_images=result['failed_images'],
            output_folder=output_folder
        )
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"分析过程中发生错误: {str(e)}")


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host=API.host, port=API.port)
